import asyncio

import neo4j
from my_extractor.openai_ere import (
    LLMEntityRelationExtractor,
    OnError,
)
from neo4j_graphrag.experimental.components.kg_writer import Neo4jWriter
from neo4j_graphrag.experimental.components.pdf_loader import PdfLoader
from neo4j_graphrag.experimental.components.resolver import (
    SinglePropertyExactMatchResolver,
)
from neo4j_graphrag.experimental.components.schema import (
    SchemaBuilder,
    NodeType,
    PropertyType,
    RelationshipType,
)
from langchain_text_splitters import RecursiveCharacterTextSplitter
from neo4j_graphrag.experimental.components.text_splitters.langchain import LangChainTextSplitterAdapter
from neo4j_graphrag.experimental.pipeline import Pipeline
from neo4j_graphrag.llm import LLMInterface, OpenAILLM

max_len = 400
overlap = 50
# 文本拆分器
sep_patterns = [
    '```\n', r'\n\*\*\*+\n', r'\n---+\n', r'\n___+\n',
    r'\n\n', r'\n', "。|！|？", r"\.\s|\!\s|\?\s", r"；|;\s"
]


async def define_and_run_pipeline(
        neo4j_driver: neo4j.Driver, llm_config
) -> None:
    """This is where we define and run the KG builder pipeline, instantiating a few
    components:
    - Text Splitter: in this example we use the fixed size text splitter
    - Schema Builder: this component takes a list of entities, relationships and
        possible triplets as inputs, validate them and return a schema ready to use
        for the rest of the pipeline
    - LLM Entity Relation Extractor is an LLM-based entity and relation extractor:
        based on the provided schema, the LLM will do its best to identity these
        entities and their relations within the provided text
    - KG writer: once entities and relations are extracted, they can be writen
        to a Neo4j database
    """
    pipe = Pipeline()
    # define the components
    pipe.add_component(PdfLoader(), "loader")
    pipe.add_component(
        LangChainTextSplitterAdapter(RecursiveCharacterTextSplitter(
            separators=sep_patterns,
            is_separator_regex=True,
            chunk_size=max_len,
            chunk_overlap=overlap,
            strip_whitespace=True
        )),
        "splitter",
    )
    pipe.add_component(SchemaBuilder(), "schema")
    pipe.add_component(
        LLMEntityRelationExtractor(
            llm_config=llm_config,
            on_error=OnError.IGNORE,
        ),
        "extractor",
    )
    pipe.add_component(Neo4jWriter(neo4j_driver), "writer")
    pipe.add_component(SinglePropertyExactMatchResolver(neo4j_driver), "resolver")
    # define the execution order of component
    # and how the output of previous components must be used
    pipe.connect("loader", "splitter", {"text": "loader.text"})
    pipe.connect("splitter", "extractor", input_config={"chunks": "splitter"})
    pipe.connect(
        "schema",
        "extractor",
        input_config={"schema": "schema", "document_info": "loader.document_info"},
    )
    pipe.connect(
        "extractor",
        "writer",
        input_config={"graph": "extractor"},
    )
    pipe.connect("writer", "resolver", {})
    # user input:
    # the initial text
    # and the list of entities and relations we are looking for
    pipe_inputs = {
        "loader": {},
        "schema": {
            "node_types": [
                # 核心文件类
                NodeType(
                    label="政策文件",
                    properties=[
                        PropertyType(name="文件名称", type="STRING"),
                        PropertyType(name="文件类型", type="STRING"),  # 条例、计划、指南等
                        PropertyType(name="发布时间", type="DATE"),
                        PropertyType(name="执行起止时间", type="STRING"),
                        PropertyType(name="内容摘要", type="STRING"),
                    ],
                ),
                # 发布相关
                NodeType(
                    label="发布部门",
                    properties=[
                        PropertyType(name="部门名称", type="STRING"),
                        PropertyType(name="级别", type="STRING"),
                    ],
                ),
                NodeType(
                    label="发布地点",
                    properties=[
                        PropertyType(name="地点名称", type="STRING"),
                        PropertyType(name="行政区划", type="STRING"),
                    ],
                ),
                # 行动计划专属
                NodeType(
                    label="战略目标",
                    properties=[
                        PropertyType(name="目标名称", type="STRING"),
                        PropertyType(name="目标描述", type="STRING"),
                        PropertyType(name="目标年份", type="STRING"),
                    ],
                ),
                NodeType(
                    label="建设方向",
                    properties=[
                        PropertyType(name="方向名称", type="STRING"),
                        PropertyType(name="关键任务", type="STRING"),
                    ],
                ),
                NodeType(
                    label="应用场景",
                    properties=[
                        PropertyType(name="场景名称", type="STRING"),
                        PropertyType(name="涉及领域", type="STRING"),  # 城市/农业/工业等
                    ],
                ),
                NodeType(
                    label="参与主体",
                    properties=[
                        PropertyType(name="主体类型", type="STRING"),  # 企业、行业组织、城市等
                        PropertyType(name="主体名称", type="STRING"),
                    ],
                ),
                NodeType(
                    label="技术能力",
                    properties=[
                        PropertyType(name="能力名称", type="STRING"),
                        PropertyType(name="涉及技术", type="STRING"),  # 区块链、隐私计算等
                    ],
                ),
                NodeType(
                    label="数据类型",
                    properties=[
                        PropertyType(name="数据类型名称", type="STRING"),
                        PropertyType(name="描述", type="STRING"),
                    ],
                ),
            ],

            "relationship_types": [
                RelationshipType(label="发布"),
                RelationshipType(label="发布于"),
                RelationshipType(label="适用于"),
                RelationshipType(label="包含目标"),
                RelationshipType(label="包含方向"),
                RelationshipType(label="涵盖场景"),
                RelationshipType(label="参与"),
                RelationshipType(label="建设能力"),
                RelationshipType(label="涉及数据类型"),
            ],

            "patterns": [
                ("发布部门", "发布", "政策文件"),
                ("发布部门", "发布于", "发布地点"),
                ("政策文件", "包含目标", "战略目标"),
                ("政策文件", "包含方向", "建设方向"),
                ("政策文件", "涵盖场景", "应用场景"),
                ("政策文件", "参与", "参与主体"),
                ("政策文件", "建设能力", "技术能力"),
                ("政策文件", "涉及数据类型", "数据类型"),
            ],
        }

        ,
    }
    # run the pipeline for each documents
    for document in [
        "files/江苏省数据局（江苏省政务服务管理办公室） 省级法规政策 江苏省数据条例（2025年1月22日江苏省第十四届人民代表大会第三次会议通过）.pdf",
        "files/江苏省数据局（江苏省政务服务管理办公室） 省级法规政策 省政府办公厅关于加快释放数据要素价值培育壮大数据产业的意见（苏政办发〔2024〕34号）.pdf",
    ]:
        pipe_inputs["loader"]["filepath"] = document
        await pipe.run(pipe_inputs)


async def main() -> None:
    openai_llm_config = {'init': {'base_url': 'http://172.16.3.116:2380/v1', 'api_key': 'EMPTY'},
                         'model': 'Qwen3-32B-AWQ', 'enable_thinking': True}

    driver = neo4j.GraphDatabase.driver(
        "bolt://172.16.3.123:7687", auth=("neo4j", "1234321ly")
    )
    await define_and_run_pipeline(driver, openai_llm_config)
    driver.close()


if __name__ == "__main__":
    asyncio.run(main())
